#!/usr/bin/env python3
"""
嵌入模型示例

这个示例展示了如何使用LangChain的嵌入功能：
1. 文本嵌入和相似度计算
2. 文档嵌入和聚类分析
3. 语义搜索和匹配
4. 嵌入模型的比较和评估
"""

import sys
import os
import numpy as np
from typing import List, Dict, Any, Optional, Tuple
from rich.console import Console
from rich.panel import Panel
from rich.table import Table
from rich.progress import Progress, TaskID

# 可选依赖
try:
    import matplotlib.pyplot as plt
    HAS_MATPLOTLIB = True
except ImportError:
    HAS_MATPLOTLIB = False

try:
    from sklearn.cluster import KMeans
    from sklearn.decomposition import PCA
    from sklearn.metrics.pairwise import cosine_similarity
    HAS_SKLEARN = True
except ImportError:
    HAS_SKLEARN = False

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from core.embedding_manager import embedding_manager

console = Console()

class EmbeddingDemo:
    """嵌入模型演示类"""
    
    def __init__(self):
        """初始化演示"""
        self.console = console
        self.sample_texts = [
            "Python是一种高级编程语言",
            "机器学习是人工智能的一个分支",
            "深度学习使用神经网络进行学习",
            "自然语言处理处理人类语言",
            "计算机视觉让机器理解图像",
            "数据科学结合统计学和编程",
            "人工智能正在改变世界",
            "算法是解决问题的步骤",
            "大数据需要特殊的处理技术",
            "云计算提供按需计算资源"
        ]
        
        self.document_categories = {
            "编程语言": [
                "Python是一种解释型、高级编程语言，具有动态语义。",
                "Java是一种面向对象的编程语言，具有平台无关性。",
                "JavaScript是一种脚本语言，主要用于网页开发。",
                "C++是一种通用的编程语言，支持面向对象编程。"
            ],
            "机器学习": [
                "监督学习使用标记数据训练模型进行预测。",
                "无监督学习从未标记数据中发现隐藏模式。",
                "强化学习通过与环境交互来学习最优策略。",
                "深度学习使用多层神经网络处理复杂数据。"
            ],
            "数据科学": [
                "数据清洗是数据科学项目的重要第一步。",
                "统计分析帮助理解数据的分布和关系。",
                "数据可视化使复杂数据变得易于理解。",
                "特征工程是提高模型性能的关键技术。"
            ]
        }
    
    def display_welcome(self):
        """显示欢迎信息"""
        welcome_text = """
🔤 [bold cyan]嵌入模型功能演示[/bold cyan]

本示例将展示：
• 文本嵌入和向量化
• 语义相似度计算
• 文档聚类分析
• 语义搜索功能
• 嵌入模型比较

让我们开始探索嵌入模型的强大功能！
        """
        
        welcome_panel = Panel(
            welcome_text,
            title="嵌入模型演示",
            border_style="cyan",
            padding=(1, 2)
        )
        
        self.console.print(welcome_panel)
    
    def check_embedding_status(self):
        """检查嵌入模型状态"""
        self.console.print("\n[bold cyan]📊 嵌入模型状态检查[/bold cyan]\n")
        
        # 获取模型信息
        model_info = embedding_manager.get_model_info()
        
        status_table = Table(title="嵌入模型状态", show_header=True, header_style="bold magenta")
        status_table.add_column("项目", style="cyan", no_wrap=True)
        status_table.add_column("值", style="green")
        
        status_table.add_row("模型就绪", "✅ 是" if embedding_manager.is_ready() else "❌ 否")
        status_table.add_row("当前模型", model_info.get('model_name', '未设置'))
        status_table.add_row("模型类型", model_info.get('model_class', '未知'))
        status_table.add_row("向量维度", str(model_info.get('dimensions', '未知')))
        
        self.console.print(status_table)
        
        # 列出可用模型
        available_models = embedding_manager.list_available_models()
        if available_models:
            models_table = Table(title="可用嵌入模型", show_header=True, header_style="bold magenta")
            models_table.add_column("模型名称", style="cyan")
            models_table.add_column("描述", style="green")
            
            for name, description in available_models.items():
                models_table.add_row(name, description)
            
            self.console.print(models_table)
    
    def demo_basic_embedding(self):
        """演示基础嵌入功能"""
        self.console.print("\n[bold cyan]🔤 基础嵌入功能演示[/bold cyan]\n")
        
        if not embedding_manager.is_ready():
            self.console.print("[red]❌ 嵌入模型未就绪[/red]")
            return
        
        # 测试单个文本嵌入
        test_text = "人工智能正在改变我们的世界"
        self.console.print(f"测试文本: [green]{test_text}[/green]")
        
        try:
            # 嵌入单个查询
            embedding = embedding_manager.embed_query(test_text)
            
            self.console.print(f"✅ 嵌入成功")
            self.console.print(f"向量维度: {len(embedding)}")
            self.console.print(f"向量前5个值: {embedding[:5]}")
            
            # 测试文档批量嵌入
            self.console.print(f"\n批量嵌入 {len(self.sample_texts)} 个文档...")
            
            embeddings = embedding_manager.embed_documents(self.sample_texts)
            
            self.console.print(f"✅ 批量嵌入成功")
            self.console.print(f"嵌入矩阵形状: {len(embeddings)} x {len(embeddings[0])}")
            
            return embeddings
            
        except Exception as e:
            self.console.print(f"[red]❌ 嵌入失败: {e}[/red]")
            return None
    
    def demo_similarity_calculation(self, embeddings: List[List[float]]):
        """演示相似度计算"""
        self.console.print("\n[bold cyan]📐 语义相似度计算演示[/bold cyan]\n")
        
        if not embeddings:
            self.console.print("[red]❌ 没有可用的嵌入向量[/red]")
            return
        
        # 计算相似度矩阵
        similarity_matrix = cosine_similarity(embeddings)
        
        # 找出最相似的文本对
        similarities = []
        for i in range(len(self.sample_texts)):
            for j in range(i + 1, len(self.sample_texts)):
                similarity = similarity_matrix[i][j]
                similarities.append((i, j, similarity))
        
        # 按相似度排序
        similarities.sort(key=lambda x: x[2], reverse=True)
        
        # 显示最相似的文本对
        similarity_table = Table(title="最相似的文本对", show_header=True, header_style="bold magenta")
        similarity_table.add_column("文本1", style="cyan", width=30)
        similarity_table.add_column("文本2", style="green", width=30)
        similarity_table.add_column("相似度", style="yellow", width=10)
        
        for i, (idx1, idx2, sim) in enumerate(similarities[:5]):
            text1 = self.sample_texts[idx1][:25] + "..." if len(self.sample_texts[idx1]) > 25 else self.sample_texts[idx1]
            text2 = self.sample_texts[idx2][:25] + "..." if len(self.sample_texts[idx2]) > 25 else self.sample_texts[idx2]
            similarity_table.add_row(text1, text2, f"{sim:.3f}")
        
        self.console.print(similarity_table)
        
        return similarity_matrix
    
    def demo_semantic_search(self):
        """演示语义搜索"""
        self.console.print("\n[bold cyan]🔍 语义搜索演示[/bold cyan]\n")
        
        if not embedding_manager.is_ready():
            self.console.print("[red]❌ 嵌入模型未就绪[/red]")
            return
        
        # 创建文档库
        all_documents = []
        doc_categories = []
        
        for category, docs in self.document_categories.items():
            all_documents.extend(docs)
            doc_categories.extend([category] * len(docs))
        
        # 嵌入所有文档
        self.console.print("正在嵌入文档库...")
        try:
            doc_embeddings = embedding_manager.embed_documents(all_documents)
            self.console.print(f"✅ 已嵌入 {len(all_documents)} 个文档")
        except Exception as e:
            self.console.print(f"[red]❌ 文档嵌入失败: {e}[/red]")
            return
        
        # 测试查询
        queries = [
            "如何学习编程？",
            "什么是神经网络？",
            "数据分析的方法有哪些？"
        ]
        
        for query in queries:
            self.console.print(f"\n查询: [green]{query}[/green]")
            
            try:
                # 嵌入查询
                query_embedding = embedding_manager.embed_query(query)
                
                # 计算与所有文档的相似度
                similarities = cosine_similarity([query_embedding], doc_embeddings)[0]
                
                # 获取最相似的文档
                top_indices = np.argsort(similarities)[::-1][:3]
                
                results_table = Table(title=f"搜索结果: {query}", show_header=True, header_style="bold magenta")
                results_table.add_column("排名", style="cyan", width=6)
                results_table.add_column("类别", style="yellow", width=12)
                results_table.add_column("文档", style="green", width=40)
                results_table.add_column("相似度", style="white", width=10)
                
                for rank, idx in enumerate(top_indices, 1):
                    doc = all_documents[idx]
                    category = doc_categories[idx]
                    similarity = similarities[idx]
                    
                    doc_preview = doc[:35] + "..." if len(doc) > 35 else doc
                    results_table.add_row(str(rank), category, doc_preview, f"{similarity:.3f}")
                
                self.console.print(results_table)
                
            except Exception as e:
                self.console.print(f"[red]❌ 搜索失败: {e}[/red]")
    
    def demo_document_clustering(self):
        """演示文档聚类"""
        self.console.print("\n[bold cyan]📊 文档聚类演示[/bold cyan]\n")
        
        if not HAS_SKLEARN:
            self.console.print("[yellow]⚠️ 需要安装scikit-learn才能运行聚类演示[/yellow]")
            self.console.print("[dim]安装命令: pip install scikit-learn[/dim]\n")
            return
        
        if not embedding_manager.is_ready():
            self.console.print("[red]❌ 嵌入模型未就绪[/red]")
            return
        
        # 准备文档
        documents = [
            "Python是一种编程语言",
            "机器学习是人工智能的分支",
            "深度学习使用神经网络",
            "Java是面向对象的编程语言",
            "自然语言处理处理文本数据",
            "计算机视觉处理图像数据",
            "JavaScript用于网页开发",
            "数据科学分析大量数据"
        ]
        
        self.console.print("正在生成文档嵌入...")
        
        # 生成嵌入
        embeddings = []
        with Progress() as progress:
            task = progress.add_task("生成嵌入...", total=len(documents))
            
            for doc in documents:
                try:
                    embedding = embedding_manager.embed_text(doc)
                    embeddings.append(embedding)
                    progress.advance(task)
                except Exception as e:
                    self.console.print(f"[red]生成嵌入失败: {e}[/red]")
                    return
        
        if not embeddings:
            self.console.print("[red]❌ 未能生成任何嵌入[/red]")
            return
        
        # 转换为numpy数组
        embeddings_array = np.array(embeddings)
        
        # 执行聚类
        n_clusters = 3
        kmeans = KMeans(n_clusters=n_clusters, random_state=42, n_init=10)
        cluster_labels = kmeans.fit_predict(embeddings_array)
        
        # 显示聚类结果
        cluster_table = Table(title="文档聚类结果", show_header=True, header_style="bold magenta")
        cluster_table.add_column("文档", style="green", width=40)
        cluster_table.add_column("聚类", style="cyan", width=10)
        
        for doc, label in zip(documents, cluster_labels):
            cluster_table.add_row(doc, f"聚类 {label}")
        
        self.console.print(cluster_table)
        
        # 分析聚类
        clusters = {}
        for i, label in enumerate(cluster_labels):
            if label not in clusters:
                clusters[label] = []
            clusters[label].append(documents[i])
        
        self.console.print(f"\n[bold green]聚类分析:[/bold green]")
        for cluster_id, docs in clusters.items():
            self.console.print(f"\n[cyan]聚类 {cluster_id}:[/cyan]")
            for doc in docs:
                self.console.print(f"  • {doc}")
    
    def demo_visualization(self):
        """演示嵌入可视化"""
        self.console.print("\n[bold cyan]📈 嵌入可视化演示[/bold cyan]\n")
        
        if not HAS_MATPLOTLIB or not HAS_SKLEARN:
            missing_deps = []
            if not HAS_MATPLOTLIB:
                missing_deps.append("matplotlib")
            if not HAS_SKLEARN:
                missing_deps.append("scikit-learn")
            
            self.console.print(f"[yellow]⚠️ 需要安装 {', '.join(missing_deps)} 才能运行可视化演示[/yellow]")
            self.console.print(f"[dim]安装命令: pip install {' '.join(missing_deps)}[/dim]\n")
            return
        
        if not embedding_manager.is_ready():
            self.console.print("[red]❌ 嵌入模型未就绪[/red]")
            return
        
        # 准备不同类别的文本
        texts = {
            "编程语言": ["Python编程", "Java开发", "JavaScript前端", "C++系统编程"],
            "机器学习": ["深度学习", "神经网络", "机器学习算法", "人工智能"],
            "数据科学": ["数据分析", "统计学习", "数据挖掘", "大数据处理"]
        }
        
        all_texts = []
        labels = []
        colors = ['red', 'blue', 'green']
        
        for i, (category, category_texts) in enumerate(texts.items()):
            all_texts.extend(category_texts)
            labels.extend([i] * len(category_texts))
        
        self.console.print("正在生成嵌入向量...")
        
        # 生成嵌入
        embeddings = []
        with Progress() as progress:
            task = progress.add_task("生成嵌入...", total=len(all_texts))
            
            for text in all_texts:
                try:
                    embedding = embedding_manager.embed_text(text)
                    embeddings.append(embedding)
                    progress.advance(task)
                except Exception as e:
                    self.console.print(f"[red]生成嵌入失败: {e}[/red]")
                    return
        
        if not embeddings:
            self.console.print("[red]❌ 未能生成任何嵌入[/red]")
            return
        
        # 降维到2D
        embeddings_array = np.array(embeddings)
        pca = PCA(n_components=2)
        embeddings_2d = pca.fit_transform(embeddings_array)
        
        # 创建可视化
        plt.figure(figsize=(10, 8))
        
        for i, category in enumerate(texts.keys()):
            mask = np.array(labels) == i
            plt.scatter(embeddings_2d[mask, 0], embeddings_2d[mask, 1], 
                       c=colors[i], label=category, alpha=0.7, s=100)
        
        # 添加文本标签
        for i, text in enumerate(all_texts):
            plt.annotate(text, (embeddings_2d[i, 0], embeddings_2d[i, 1]), 
                        xytext=(5, 5), textcoords='offset points', 
                        fontsize=8, alpha=0.8)
        
        plt.title('文本嵌入2D可视化', fontsize=16, fontweight='bold')
        plt.xlabel(f'主成分1 (解释方差: {pca.explained_variance_ratio_[0]:.2%})')
        plt.ylabel(f'主成分2 (解释方差: {pca.explained_variance_ratio_[1]:.2%})')
        plt.legend()
        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        
        # 保存图片
        plt.savefig('embedding_visualization.png', dpi=300, bbox_inches='tight')
        self.console.print("[green]✅ 可视化图片已保存为 'embedding_visualization.png'[/green]")
        
        # 显示图片（如果在支持的环境中）
        try:
            plt.show()
        except Exception:
            self.console.print("[yellow]无法显示图片，但已保存到文件[/yellow]")
        
        plt.close()
    
    def demo_model_comparison(self):
        """演示不同嵌入模型的比较"""
        self.console.print("\n[bold cyan]⚖️ 嵌入模型比较演示[/bold cyan]\n")
        
        available_models = embedding_manager.list_available_models()
        if len(available_models) < 2:
            self.console.print("[yellow]⚠️ 需要至少2个可用模型进行比较[/yellow]")
            return
        
        # 选择要比较的模型
        model_names = list(available_models.keys())[:2]  # 取前两个模型
        
        test_queries = [
            "机器学习算法",
            "Python编程语言",
            "数据科学分析"
        ]
        
        comparison_results = {}
        
        for model_name in model_names:
            self.console.print(f"正在测试模型: {model_name}")
            
            # 切换模型
            if embedding_manager.set_model(model_name):
                try:
                    # 测试嵌入性能
                    embeddings = embedding_manager.embed_documents(test_queries)
                    
                    # 计算相似度
                    similarities = cosine_similarity(embeddings)
                    avg_similarity = np.mean(similarities[np.triu_indices_from(similarities, k=1)])
                    
                    comparison_results[model_name] = {
                        'avg_similarity': avg_similarity,
                        'vector_dim': len(embeddings[0]),
                        'status': '✅ 成功'
                    }
                    
                except Exception as e:
                    comparison_results[model_name] = {
                        'avg_similarity': 0,
                        'vector_dim': 0,
                        'status': f'❌ 失败: {str(e)[:30]}'
                    }
            else:
                comparison_results[model_name] = {
                    'avg_similarity': 0,
                    'vector_dim': 0,
                    'status': '❌ 切换失败'
                }
        
        # 显示比较结果
        comparison_table = Table(title="模型比较结果", show_header=True, header_style="bold magenta")
        comparison_table.add_column("模型名称", style="cyan", width=25)
        comparison_table.add_column("向量维度", style="green", width=12)
        comparison_table.add_column("平均相似度", style="yellow", width=15)
        comparison_table.add_column("状态", style="white", width=20)
        
        for model_name, results in comparison_results.items():
            comparison_table.add_row(
                model_name,
                str(results['vector_dim']),
                f"{results['avg_similarity']:.3f}",
                results['status']
            )
        
        self.console.print(comparison_table)
    
    def demo_embedding_applications(self):
        """演示嵌入的实际应用场景"""
        self.console.print("\n[bold cyan]🚀 嵌入应用场景演示[/bold cyan]\n")
        
        applications = [
            {
                "name": "问答匹配",
                "description": "根据问题找到最相关的答案",
                "demo_func": self._demo_qa_matching
            },
            {
                "name": "文档去重",
                "description": "识别和移除重复或相似的文档",
                "demo_func": self._demo_document_deduplication
            },
            {
                "name": "内容推荐",
                "description": "基于用户兴趣推荐相关内容",
                "demo_func": self._demo_content_recommendation
            }
        ]
        
        for app in applications:
            self.console.print(f"\n[bold green]📋 {app['name']}[/bold green]")
            self.console.print(f"[dim]{app['description']}[/dim]")
            
            try:
                app['demo_func']()
            except Exception as e:
                self.console.print(f"[red]❌ 演示失败: {e}[/red]")
    
    def _demo_qa_matching(self):
        """演示问答匹配"""
        qa_pairs = [
            ("什么是Python？", "Python是一种高级编程语言，以其简洁和可读性著称。"),
            ("如何学习机器学习？", "学习机器学习需要掌握数学基础、编程技能和实践经验。"),
            ("数据科学的工具有哪些？", "常用的数据科学工具包括Python、R、SQL、Tableau等。"),
            ("什么是深度学习？", "深度学习是机器学习的一个分支，使用多层神经网络。")
        ]
        
        questions = [q for q, _ in qa_pairs]
        answers = [a for _, a in qa_pairs]
        
        # 测试问题
        test_question = "Python编程语言是什么？"
        
        try:
            # 嵌入问题和答案
            question_embeddings = embedding_manager.embed_documents(questions)
            test_embedding = embedding_manager.embed_query(test_question)
            
            # 找到最匹配的问题
            similarities = cosine_similarity([test_embedding], question_embeddings)[0]
            best_match_idx = np.argmax(similarities)
            
            self.console.print(f"测试问题: [green]{test_question}[/green]")
            self.console.print(f"最匹配问题: [cyan]{questions[best_match_idx]}[/cyan]")
            self.console.print(f"对应答案: [yellow]{answers[best_match_idx]}[/yellow]")
            self.console.print(f"相似度: {similarities[best_match_idx]:.3f}")
            
        except Exception as e:
            self.console.print(f"[red]问答匹配失败: {e}[/red]")
    
    def _demo_document_deduplication(self):
        """演示文档去重"""
        documents = [
            "Python是一种编程语言",
            "Python是一种高级编程语言",  # 相似文档
            "机器学习是AI的分支",
            "Java是面向对象的语言",
            "机器学习属于人工智能领域",  # 相似文档
            "数据科学结合统计和编程"
        ]
        
        try:
            # 嵌入所有文档
            embeddings = embedding_manager.embed_documents(documents)
            
            # 计算相似度矩阵
            similarity_matrix = cosine_similarity(embeddings)
            
            # 找出重复文档（相似度阈值0.8）
            threshold = 0.8
            duplicates = []
            
            for i in range(len(documents)):
                for j in range(i + 1, len(documents)):
                    if similarity_matrix[i][j] > threshold:
                        duplicates.append((i, j, similarity_matrix[i][j]))
            
            self.console.print("文档去重分析:")
            if duplicates:
                for i, j, sim in duplicates:
                    self.console.print(f"  相似文档对 (相似度: {sim:.3f}):")
                    self.console.print(f"    文档1: {documents[i]}")
                    self.console.print(f"    文档2: {documents[j]}")
            else:
                self.console.print("  未发现重复文档")
                
        except Exception as e:
            self.console.print(f"[red]文档去重失败: {e}[/red]")
    
    def _demo_content_recommendation(self):
        """演示内容推荐"""
        user_interests = ["机器学习", "Python编程"]
        content_library = [
            "深度学习神经网络教程",
            "Python数据分析指南",
            "Java Web开发实战",
            "机器学习算法详解",
            "前端JavaScript框架",
            "Python机器学习库介绍",
            "数据库设计原理",
            "人工智能伦理讨论"
        ]
        
        try:
            # 嵌入用户兴趣和内容库
            interest_embeddings = embedding_manager.embed_documents(user_interests)
            content_embeddings = embedding_manager.embed_documents(content_library)
            
            # 计算用户兴趣的平均嵌入
            avg_interest_embedding = np.mean(interest_embeddings, axis=0)
            
            # 计算与内容的相似度
            similarities = cosine_similarity([avg_interest_embedding], content_embeddings)[0]
            
            # 排序推荐
            recommendations = sorted(
                zip(content_library, similarities),
                key=lambda x: x[1],
                reverse=True
            )
            
            self.console.print(f"用户兴趣: {', '.join(user_interests)}")
            self.console.print("推荐内容:")
            
            for i, (content, sim) in enumerate(recommendations[:5], 1):
                self.console.print(f"  {i}. {content} (相关度: {sim:.3f})")
                
        except Exception as e:
            self.console.print(f"[red]内容推荐失败: {e}[/red]")
    
    def run_demo(self):
        """运行完整演示"""
        self.display_welcome()
        
        # 检查嵌入模型状态
        self.check_embedding_status()
        
        if not embedding_manager.is_ready():
            self.console.print("\n[red]❌ 嵌入模型未就绪，无法继续演示[/red]")
            self.console.print("[yellow]请检查您的配置和API密钥[/yellow]")
            return
        
        # 运行各种演示
        embeddings = self.demo_basic_embedding()
        
        if embeddings:
            self.demo_similarity_calculation(embeddings)
        
        self.demo_semantic_search()
        self.demo_document_clustering()
        self.demo_visualization()
        self.demo_model_comparison()
        self.demo_embedding_applications()
        
        # 结束信息
        self.console.print("\n[bold green]🎉 嵌入模型演示完成！[/bold green]")
        self.console.print("\n[cyan]通过这个演示，您已经了解了：[/cyan]")
        self.console.print("• 如何使用嵌入模型进行文本向量化")
        self.console.print("• 语义相似度计算和搜索")
        self.console.print("• 文档聚类和分类")
        self.console.print("• 嵌入模型的实际应用场景")

def main():
    """主函数"""
    try:
        demo = EmbeddingDemo()
        demo.run_demo()
    
    except KeyboardInterrupt:
        console.print("\n[cyan]演示已中断[/cyan]")
    
    except Exception as e:
        console.print(f"[red]演示失败: {str(e)}[/red]")

if __name__ == "__main__":
    main() 